from typing import Optional

import pandas as pd

from model.action import Action
from model.networkstate import NetworkState


class NetworkTrace:
    def __init__(self):
        self.len = 0
        self.time = 0
        self.time_trace = []
        self.state_trace = []
        self.act_trace = []

    def __len__(self):
        return self.len

    def __getitem__(self, item):
        return  self.state_trace[item]

    def __repr__(self):
        return str(self.to_component_state_df())

    def append(self, time, state: NetworkState, act: Optional[Action]):
        d = time - self.time
        self.time = time
        self.len += 1
        self.time_trace.append(d)
        self.state_trace.append(state)
        if not act is None:
            self.act_trace.append(act)
        else:
            self.act_trace.append(None)


    def to_network_state_df(self):
        df = pd.DataFrame({'Time': self.time_trace, 'State': self.state_trace, 'Action': self.act_trace})
        return df

    def to_component_state_df(self):
        if len(self.state_trace) == 0:
            return None
        else:
            components = self.state_trace[0].keys()
            df = pd.DataFrame({'Time': self.time_trace}).astype(float)
            for component in components:
                new_col = pd.DataFrame({str(component): [{str(k):str(v) for k,v in state[component].items()} for state in self.state_trace]}).astype(str)
                df = df.join(new_col, how="left")
            new_col = pd.DataFrame({'Action' : self.act_trace}).astype(str)
            df = df.join(new_col, how="left")
            return df

    def to_puml(self):
        #  todo 完成画图模块
        pass




if __name__ == '__main__':
    from nptem import NPTEM
    import timeit

    print(timeit.timeit(NPTEM.mlag_run, number=1))
    nt = NPTEM.mlag_run()
    df_ns = nt.to_network_state_df()
    df_cs = nt.to_component_state_df()
    print(df_cs)